#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File     : datasets.py
@Project  : pipecoco
@Date     : 2021/8/23
@Author   : Zhang Jinyang
@Contact  : zhang-jy@sjtu.edu.cn
'''

import mindspore.dataset as de
import mindspore.dataset.transforms.c_transforms as C
import mindspore.dataset.vision.c_transforms as CV
from PIL import ImageFile
from mindspore.common import dtype as mstype

ImageFile.LOAD_TRUNCATED_IMAGES = True

def create_dataset_cifar10(data_dir, image_size, batch_size):

    """
    create dataset for train or test
    """
    cifar_ds = de.Cifar10Dataset(data_dir,shuffle=False)

    resize_op = CV.Resize(image_size)
    rescale_op = CV.Rescale(rescale=1.0 / 255.0, shift=0.0)
    normalize_op = CV.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))

    channel_swap_op = CV.HWC2CHW()
    typecast_op = C.TypeCast(mstype.int32)
    cifar_ds = cifar_ds.map(input_columns="label", operations=typecast_op, num_parallel_workers=4)

    cifar_ds = cifar_ds.map(input_columns="image", operations=resize_op, num_parallel_workers=4)
    cifar_ds = cifar_ds.map(input_columns="image", operations=rescale_op, num_parallel_workers=4)
    cifar_ds = cifar_ds.map(input_columns="image", operations=normalize_op, num_parallel_workers=4)
    cifar_ds = cifar_ds.map(input_columns="image", operations=channel_swap_op, num_parallel_workers=4)

    cifar_ds = cifar_ds.batch(batch_size, drop_remainder=True)

    return cifar_ds